# uncompyle6 version 3.2.3
# Python bytecode 3.6 (3379)
# Decompiled from: Python 3.6.8 |Anaconda custom (64-bit)| (default, Feb 21 2019, 18:30:04) [MSC v.1916 64 bit (AMD64)]
# Embedded file name: reporter\generators\ds_target.py
import json, pystache, logging
from urllib.parse import urlparse
from sqlalchemy import select as sql_select
from sqlalchemy import func as sql_func
from sqlalchemy import or_ as sql_or
from sqlalchemy import and_ as sql_and
from sqlalchemy import desc as sql_desc
from sqlalchemy import distinct as sql_distinct
from reporter.generators.generic import DataSource
from db import Connection
from db.tables.target_vulns import TargetVulnRow, TargetVulnsTable
from db.tables.targets import TargetRow, TargetsTable, TargetTypes
from db.tables.users import UserRow, UsersTable
from db.tables.vuln_types import VulnTypeRow, VulnTypesTable
from helpers.address import extract_base_url_from_address

logger = logging.getLogger("reporter.generators.ds_target")
vuln_type_filters = (
    "cvss3",
    "long_description",
    "name",
    "recommendation",
    "cvss_score",
    "description",
    "cvss2",
    "vt_id",
    "severity",
    "impact",
    "app_id",
)
vuln_filter = (
    "vt_id",
    "vuln_hash",
    "request",
    "source",
    "details",
    "app_id",
    "loc_detail",
    "request_secure",
    "request_port",
    "attack_vector",
    "use_ssl",
    "first_seen",
    "status",
)


class DSTargetResult(DataSource):
    def __init__(self, shard_id, target_id, vuln_id_list=None, source_type=None):
        super().__init__()
        self.shard_id = shard_id
        self.target_id = target_id
        self.vuln_id_list = vuln_id_list
        self.source_type = source_type
        self.hide_web_scan_details = False
        with Connection(self.shard_id) as (db_connection):
            sf = TargetsTable.join(UsersTable, UserRow.user_id == TargetRow.owner_id)
            q = (
                sql_select((TargetRow, UserRow.license_type))
                .select_from(sf)
                .where(TargetRow.target_id == target_id)
            )
            t = db_connection.execute(q).fetchone()
            if not t:
                raise Exception("target %s not found" % (target_id,))
            self.address = t.address
            self.host = extract_base_url_from_address(self.address)
            if t["license_type"] == "trial":
                if t["type"] != TargetTypes.DEMO:
                    self.hide_web_scan_details = True

    def info(self):
        """
        required fields: host, start_url
        :return:
        """
        result = dict(
            host=self.host,
            start_url=self.address,
            vuln_selection=self.vuln_id_list is not None,
            source_type=self.source_type,
        )
        return result

    def locations(self):
        """
        required fields: {url}
        :return:
        """
        loc_query = (
            sql_select((sql_distinct(TargetVulnRow.url),))
            .where(TargetVulnRow.target_id == self.target_id)
            .order_by(TargetVulnRow.url)
        )
        if self.vuln_id_list:
            loc_query = loc_query.where(TargetVulnRow.vuln_id.in_(self.vuln_id_list))
        cursor = None
        while True:
            locations = []
            if cursor is None:
                q = loc_query
            else:
                q = loc_query.where(TargetVulnRow.url > cursor)
            with Connection(self.shard_id) as (db_connection):
                for location in db_connection.execute(q).fetchall():
                    location_dict = dict(url=location.url, loc_id=location.url, path="")
                    try:
                        parsed_url = urlparse(location.url)
                        location_dict["path"] = parsed_url.path
                    except Exception as e:
                        logger.exception('parsing "%s" failed with %s', location.url, e)

                    locations.append(location_dict)
                    cursor = location.url

            if not locations:
                break
            for location in locations:
                yield location

    def vulnerability_types(self):
        """
        required fields: (vuln_type, count)
        vuln_type: cvss3, long_description, name, recommendation, cvss_score, description, cvss, tags,
                   references, severity, impact, app_id
        :return:
        """
        select_from = VulnTypesTable.join(
            TargetVulnsTable, TargetVulnRow.vt_id == VulnTypeRow.vt_id
        )
        vt_query = sql_select(
            (VulnTypeRow, sql_func.count(VulnTypeRow.vt_id).label("vuln_count"))
        ).select_from(select_from)
        vt_query = (
            vt_query.where(TargetVulnRow.target_id == self.target_id)
            .group_by(VulnTypeRow.vt_id)
            .order_by(
                sql_desc(VulnTypeRow.severity), VulnTypeRow.name, VulnTypeRow.vt_id
            )
            .limit(50)
        )
        if self.vuln_id_list is None:
            vt_query = vt_query.where(TargetVulnRow.status == "open")
        if self.vuln_id_list:
            vt_query = vt_query.where(TargetVulnRow.vuln_id.in_(self.vuln_id_list))
        cursor = None
        while True:
            if cursor is None:
                q = vt_query
            else:
                q = vt_query.where(
                    sql_or(
                        VulnTypeRow.severity < cursor[0],
                        sql_and(
                            VulnTypeRow.severity == cursor[0],
                            sql_or(
                                VulnTypeRow.name > cursor[1],
                                sql_and(
                                    VulnTypeRow.name == cursor[1],
                                    VulnTypeRow.vt_id > cursor[2],
                                ),
                            ),
                        ),
                    )
                )
            vtypes = []
            with Connection(self.shard_id) as (db_connection):
                for vt in db_connection.execute(q).fetchall():
                    cursor = (vt.severity, vt.name, vt.vt_id)
                    vt = dict(vt)
                    refs = []
                    if vt["refs"]:
                        for ref in vt["refs"]:
                            refs.append(dict(name=ref[0], url=ref[1]))

                    vt["refs"] = refs
                    vt["type"] = self.extract_vuln_type_from_tags(vt["tags"])
                    vtypes.append((vt, vt["vuln_count"]))

                if not vtypes:
                    break
                for vtype in vtypes:
                    yield vtype

    def vulnerabilities(self, vtype=None, location=None, tags=None, severity=None):
        """
        tags is not exact match filter, is overlap filter !!!
        
        required fields: (vuln, location, vuln_type)
        vuln: loc_url, host, vt_id, vuln_hash, request, source, loc_id, details, app_id, loc_detail
        location: loc_type, source_id, root_id, parent_id, name, tags, input_data, url, path, loc_id
        vuln_type: cvss3, long_description, name, recommendation, cvss_score, description, cvss, tags, vt_id,
                   references, severity, impact, app_id
        :param vtype:
        :param location:
        :param tags:
        :param severity:
        :return:
        """
        select_from = TargetVulnsTable.join(
            VulnTypesTable, VulnTypeRow.vt_id == TargetVulnRow.vt_id
        )
        if tags:
            tags = list(tags)
        query = (
            sql_select(
                (
                    TargetVulnRow,
                    VulnTypeRow.details_template,
                    VulnTypeRow.app_id,
                    VulnTypeRow.tags.label("vt_tags"),
                    VulnTypeRow.cvss2,
                    VulnTypeRow.cvss3,
                    VulnTypeRow.cvss_score,
                    VulnTypeRow.impact,
                    VulnTypeRow.description,
                    VulnTypeRow.recommendation,
                    VulnTypeRow.long_description,
                    VulnTypeRow.refs,
                )
            )
            .select_from(select_from)
            .where(TargetVulnRow.target_id == self.target_id)
        )
        if self.vuln_id_list is None:
            query = query.where(TargetVulnRow.status == "open")
        if self.vuln_id_list:
            query = query.where(TargetVulnRow.vuln_id.in_(self.vuln_id_list))
        if vtype:
            query = query.where(TargetVulnRow.vt_id == vtype["vt_id"])
        if location:
            query = query.where(TargetVulnRow.url == location["url"])
        if severity is not None:
            query = query.where(TargetVulnRow.severity == severity)
        if tags:
            query = query.where(VulnTypeRow.tags.overlap(tags))
        query = query.order_by(
            sql_desc(TargetVulnRow.severity),
            TargetVulnRow.name,
            TargetVulnRow.url,
            TargetVulnRow.loc_detail,
            TargetVulnRow.vuln_id,
        )
        query = query.limit(50)
        cursor = None
        while True:
            if cursor is None:
                q = query
            else:
                q = query.where(
                    sql_or(
                        TargetVulnRow.severity < cursor[0],
                        sql_and(
                            TargetVulnRow.severity == cursor[0],
                            sql_or(
                                TargetVulnRow.name > cursor[1],
                                sql_and(
                                    TargetVulnRow.name == cursor[1],
                                    sql_or(
                                        TargetVulnRow.url > cursor[2],
                                        sql_and(
                                            TargetVulnRow.url == cursor[2],
                                            sql_or(
                                                TargetVulnRow.loc_detail > cursor[3],
                                                sql_and(
                                                    TargetVulnRow.loc_detail
                                                    == cursor[3],
                                                    TargetVulnRow.vuln_id > cursor[4],
                                                ),
                                            ),
                                        ),
                                    ),
                                ),
                            ),
                        ),
                    )
                )
            vulns = []
            with Connection(self.shard_id) as (db_connection):
                for vuln in db_connection.execute(q).fetchall():
                    cursor = (
                        vuln.severity,
                        vuln.name,
                        vuln.url,
                        vuln.loc_detail,
                        vuln.vuln_id,
                    )
                    data = dict(vuln)
                    try:
                        if data["details_type"] == "json":
                            data["details"] = pystache.render(
                                data["details_template"], json.loads(data["details"])
                            )
                    except Exception as e:
                        logger.exception(
                            "render_details failed with %s for %s", e, vuln
                        )

                    vuln = {key: data[key] for key in vuln_filter}
                    vuln["host"] = self.host
                    vuln["loc_url"] = data["url"]
                    location = dict(url=data["url"], loc_id=data["url"], path="")
                    try:
                        parsed_url = urlparse(data["url"])
                        location["path"] = parsed_url.path
                    except Exception as e:
                        logger.exception('parsing "%s" failed with %s', data["url"], e)

                    vuln_type = {key: data[key] for key in vuln_type_filters}
                    vuln_type["tags"] = data["vt_tags"]
                    refs = []
                    if data["refs"]:
                        for ref in data["refs"]:
                            refs.append(dict(name=ref[0], url=ref[1]))

                    vuln_type["references"] = refs
                    vuln_type["type"] = self.extract_vuln_type_from_tags(
                        vuln_type["tags"]
                    )
                    vulns.append((vuln, location, vuln_type))

            if not vulns:
                break
            for vuln in vulns:
                yield vuln
